# coding:utf-8
from __future__ import print_function
import os
import json
import base64
import uuid
from volcengine.visual.VisualService import VisualService
import requests

# 固定提示词
prompt = "画一张猫吃饭的图片"
print(f"使用固定提示词: {prompt}")

if __name__ == '__main__':
    visual_service = VisualService()

    # 从环境变量获取密钥
    access_key = os.getenv('VOLCANO_ACCESS_KEY')
    secret_key = os.getenv('VOLCANO_SECRET_KEY')

    if not access_key or not secret_key:
        print("错误: 请先设置环境变量VOLCANO_ACCESS_KEY和VOLCANO_SECRET_KEY")
        exit(1)

    # 设置访问密钥
    visual_service.set_ak(access_key)
    visual_service.set_sk(secret_key)

    # 构建请求参数
    form = {
        "req_key": "high_aes_general_v21_L",
        "prompt": prompt,
        "seed": -1,
        "scale": 3.5,
        "ddim_steps": 25,
        "width": 512,
        "height": 512,
        "use_pre_llm": True,
        "use_sr": True,
        "return_url": True,
        "logo_info": {
            "add_logo": False,
            "position": 0,
            "language": 0,
            "opacity": 0.3,
            "logo_text_content": ""
        }
    }

    print("----- 生成图片中，请稍候 -----")

    # 调用API
    result = visual_service.cv_process(form)

    # 处理响应
    if isinstance(result, dict) and result.get("code") == 10000:
        print("图片生成成功！")

        # 显示返回的图片URL
        if "data" in result and "image_urls" in result["data"] and result["data"]["image_urls"]:
            print("图片URL:")
            for url in result["data"]["image_urls"]:
                print(url)

            # 询问是否保存图片
            save_image = input("是否保存图片到本地？(y/n): ")
            if save_image.lower() == 'y':
                # 创建不使用代理的会话
                session = requests.Session()
                session.trust_env = False  # 不使用环境变量中的代理设置

                for i, img_url in enumerate(result["data"]["image_urls"]):
                    try:
                        print(f"正在下载图片: {img_url}")
                        img_response = session.get(img_url, timeout=30)
                        if img_response.status_code == 200:
                            filename = f"generated_image_{i}_{uuid.uuid4().hex[:8]}.jpg"
                            with open(filename, "wb") as f:
                                f.write(img_response.content)
                            print(f"图片已保存为: {filename}")
                        else:
                            print(f"下载图片失败: 状态码 {img_response.status_code}")
                    except Exception as e:
                        print(f"下载图片时出错: {str(e)}")
                        print("尝试手动从浏览器下载图片URL")

        # 如果API返回了base64编码的图片数据
        if "data" in result and "binary_data_base64" in result["data"] and result["data"]["binary_data_base64"]:
            print("获取到Base64图片数据")
            for i, base64_data in enumerate(result["data"]["binary_data_base64"]):
                try:
                    filename = f"generated_image_base64_{i}_{uuid.uuid4().hex[:8]}.jpg"
                    with open(filename, "wb") as f:
                        f.write(base64.b64decode(base64_data))
                    print(f"Base64图片已保存为: {filename}")
                except Exception as e:
                    print(f"保存Base64图片时出错: {str(e)}")

        # 显示扩写后的提示词
        if "data" in result and "rephraser_result" in result["data"] and result["data"]["rephraser_result"]:
            print("\n扩写后的提示词:")
            print(result["data"]["rephraser_result"])
    else:
        print(f"生成失败: {result.get('message', '未知错误')}")
        print(f"错误代码: {result.get('code', '未知')}")
        print(f"完整响应: {result}")